from matplotlib_inline.backend_inline import set_matplotlib_formats
import sys
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib
matplotlib.use('TkAgg')  # 或者 'Qt5Agg'
import matplotlib.pyplot as plt
from torch import nn


def use_svg_display():
    # 使用矢量图显示
    #set_matplotlib_formats('svg')  # 正确的方法
    matplotlib.use('TkAgg')  # 或者 'Qt5Agg'

def set_figsize(figsize=(3.5, 2.5)):
    use_svg_display()
    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize


#定义模型
def linreg(X, w, b): # 本函数已保存在d2lzh_pytorch包中⽅便以后使⽤
    return torch.mm(X, w) + b


#定义损失函数
def squared_loss(y_hat, y): # 本函数已保存在d2lzh_pytorch包中⽅便以后使⽤
    # 注意这⾥返回的是向量, 另外, pytorch⾥的MSELoss并没有除以 2
    return (y_hat - y.view(y_hat.size())) ** 2 / 2

#定义优化算法
def sgd(params, lr, batch_size): # 本函数已保存在d2lzh_pytorch包中⽅便以后使⽤
    for param in params:
        param.data -= lr * param.grad / batch_size # 注意这⾥更改param时⽤的param.data


#将数值标签转换为文本标签
def get_fashion_mnist_labels(labels):
    text_labels = ['t-shirt', 'trouser', 'pullover', 'dress', 'coat','sandal', 'shirt', 'sneaker', 'bag', 'ankle boot']
    return [text_labels[int(i)] for i in labels]


# 本函数已保存在d2lzh包中⽅便以后使⽤
def show_fashion_mnist(images, labels):
    # 使用合适的显示后端
    use_svg_display()  # 如果需要矢量图，保持此行，或改为交互式后端显示

    # 创建子图并设置图形大小
    _, figs = plt.subplots(1, len(images), figsize=(12, 12))

    # 遍历每个子图，显示图像及标签
    for f, img, lbl in zip(figs, images, labels):
        # 显示28x28的图片，并设置标题为标签
        f.imshow(img.view((28, 28)).numpy(), cmap='gray')  # 显示为灰度图
        f.set_title(lbl)  # 设置标题
        f.axes.get_xaxis().set_visible(False)  # 隐藏x轴
        f.axes.get_yaxis().set_visible(False)  # 隐藏y轴

    plt.show()  # 显示图形


# 定义函数来加载数据集
def load_data_fashion_mnist(batch_size, root='~/Datasets/FashionMNIST'):
    # 定义数据转换过程：将图片转换为Tensor并归一化
    transform = transforms.Compose([
        transforms.ToTensor(),  # 将PIL图片或NumPy数组转换为Tensor
        transforms.Normalize((0.5,), (0.5,))  # 归一化，均值0.5，标准差0.5
    ])

    # 加载训练集和测试集
    mnist_train = datasets.FashionMNIST(root=root, train=True, download=True, transform=transform)
    mnist_test = datasets.FashionMNIST(root=root, train=False, download=True, transform=transform)

    # 设置DataLoader来批量加载数据
    if sys.platform.startswith('win'):
        num_workers = 0  # Windows不使用额外进程加载数据
    else:
        num_workers = 4  # Linux/Mac使用4个进程加载数据

    # 创建训练集和测试集的DataLoader
    train_iter = DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    test_iter = DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=num_workers)

    return train_iter, test_iter

#评价模型在数据集上的准确率
def evaluate_accuracy(data_iter, net):
    acc_sum, n = 0.0, 0
    for X, y in data_iter:
        acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
        n += y.shape[0]
    return acc_sum / n

#softmax模型的训练过程
def train_softmax(net, train_iter, test_iter, loss, num_epochs,batch_size,params=None, lr=None, optimizer=None):
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for X, y in train_iter:
            y_hat = net(X)
            l = loss(y_hat, y).sum()

            # 梯度清零
            if optimizer is not None:
                optimizer.zero_grad()
            elif params is not None and params[0].grad is not None:
                for param in params:
                    param.grad.data.zero_()

            l.backward()
            if optimizer is None:
                sgd(params, lr, batch_size)
            else:
                optimizer.step()  # “softmax回归的简洁实现”⼀节将⽤到

            train_l_sum += l.item()
            train_acc_sum += (y_hat.argmax(dim=1) ==y).sum().item()
            n += y.shape[0]
        test_acc = evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f'% (epoch + 1, train_l_sum / n, train_acc_sum / n,test_acc))

# 本函数已保存在d2lzh_pytorch包中⽅便以后使⽤
class FlattenLayer(nn.Module):
    def __init__(self):
        super(FlattenLayer, self).__init__()

    def forward(self, x):  # x shape: (batch, *, *, ...)
        return x.view(x.shape[0], -1)


# 本函数已保存在d2lzh_pytorch包中⽅便以后使⽤
def semilogy(x_vals, y_vals, x_label, y_label, x2_vals=None, y2_vals=None,legend=None, figsize=(3.5, 2.5)):
    set_figsize(figsize)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.semilogy(x_vals, y_vals)
    if x2_vals and y2_vals:
        plt.semilogy(x2_vals, y2_vals, linestyle=':')
        plt.legend(legend)